import numpy as np  
import open3d as o3d
import os
from scipy.interpolate import griddata
from sklearn.neighbors import NearestNeighbors
import matplotlib.pyplot as plt

class DEM:
    def __init__(self, origin=np.array([0.0, 0.0, 0.0])):
        """
        初始化DEM类。

        :param origin: DEM的原点（几何中心），默认为[0.0, 0.0, 0.0]
        """
        self.Origin_coor_ = origin  # DEM的原点（几何中心）
        self.resolution_ = 1.0  # 网格分辨率
        self.width_ = 0  # DEM的宽度（theta方向）
        self.height_ = 0  # DEM的高度（phi方向）
        self.Map_elevation = {}  # 高程映射
        self.Map_color = {}  # 颜色映射
        self.Map_grad = {}  # 梯度映射

    def GetOc(self):
        """
        获取DEM的原点坐标。

        :return: DEM原点坐标（numpy数组）
        """
        return self.Origin_coor_

    def GetMap_e(self):
        """
        获取高程映射。

        :return: 高程映射（字典）
        """
        return self.Map_elevation

    def GetMap_c(self):
        """
        获取颜色映射。

        :return: 颜色映射（字典）
        """
        return self.Map_color

    def Grad_cal(self, elevation_map):
        """
        计算DEM每个网格的高程梯度。

        :param elevation_map: 高程映射（字典）
        """
        for hash_key, d_p in elevation_map.items():
            row, col = hash_key
            # 初始化相邻点的高程值
            d_a = d_b = d_c = d_d = d_e = d_f = d_g = d_h = 0
            # 定义相邻网格的键
            neighbors = {
                (row + 1, col): 'd_f',
                (row + 1, col + 1): 'd_e',
                (row + 1, col - 1): 'd_g',
                (row - 1, col): 'd_b',
                (row - 1, col + 1): 'd_c',
                (row - 1, col - 1): 'd_a',
                (row, col + 1): 'd_d',
                (row, col - 1): 'd_h',
            }
            # 更新相邻点的高程
            for neighbor, var_name in neighbors.items():
                if neighbor in elevation_map:
                    locals()[var_name] = elevation_map[neighbor]

            # 计算梯度值
            dae = abs(d_p - (d_a + d_e) * 0.5)
            dbf = abs(d_p - (d_b + d_f) * 0.5)
            dcg = abs(d_p - (d_c + d_g) * 0.5)
            ddh = abs(d_p - (d_d + d_h) * 0.5)

            grad_eva = (dae + dbf + dcg + ddh) * 0.25  # 高程梯度值
            # 存储梯度
            self.Map_grad[hash_key] = grad_eva

    def DEM_generate_polar(self, pc, gridsize):
        """
        生成极坐标下的数字高程模型 (DEM)。

        :param pc: 点云对象 (Open3D PointCloud)
        :param gridsize: 网格分辨率
        """
        self.resolution_ = gridsize
        origin = self.Origin_coor_
        print(f"DEM Origin (Cartesian): {origin}")

        # 设置极角和方位角范围
        self.height_ = np.pi  # phi [0, pi]
        self.width_ = 2 * np.pi  # theta [0, 2pi]
        nrow = int(np.ceil(self.width_ / gridsize))
        ncol = int(np.ceil(self.height_ / gridsize))

        self.width_ = nrow * gridsize
        self.height_ = ncol * gridsize

        print(f"Generating DEM with gridsize={gridsize}")
        print(f"Number of rows (theta direction): {nrow}")
        print(f"Number of columns (phi direction): {ncol}")

        # 将点从直角坐标系转变为极坐标系
        points = np.asarray(pc.points)
        colors = np.asarray(pc.colors)
        # 以DEM的原点为参考点进行平移
        translated_points = points - origin  # 平移
        x = translated_points[:, 0]
        y = translated_points[:, 1]
        z = translated_points[:, 2]
        r = np.sqrt(x**2 + y**2 + z**2)
        theta = np.arctan2(y, x)
        theta[theta < 0] += 2 * np.pi
        # 使用 np.clip 处理数值稳定性
        phi = np.arccos(np.clip(z / r, -1.0, 1.0))
        phi[r == 0] = 0

        # 计算网格索引
        rowIDs = np.floor(theta / gridsize).astype(int)
        colIDs = np.floor(phi / gridsize).astype(int)

        # 过滤超出范围的索引
        rowIDs = np.clip(rowIDs, 0, nrow - 1)
        colIDs = np.clip(colIDs, 0, ncol - 1)

        # 使用字典存储网格点索引
        pointsinGrid_id = {}
        for idx, (row, col) in enumerate(zip(rowIDs, colIDs)):
            hash_key = (row, col)
            if hash_key not in pointsinGrid_id:
                pointsinGrid_id[hash_key] = []
            pointsinGrid_id[hash_key].append(idx)

        print(f"Total unique grids with points: {len(pointsinGrid_id)}")

        # 遍历网格，计算每个网格的平均高程和颜色
        for i in range(nrow):
            for j in range(ncol):
                hash_key = (i, j)
                indices = pointsinGrid_id.get(hash_key, [])
                if len(indices) > 0:
                    average_elevation = np.mean(r[indices])
                    average_color = np.mean(colors[indices], axis=0)
                    self.Map_elevation[hash_key] = average_elevation
                    self.Map_color[hash_key] = average_color

        print(f"DEM generation completed. Total grids with elevation: {len(self.Map_elevation)}")

        # 计算高程梯度
        self.Grad_cal(self.Map_elevation)
        print(f"Gradient calculation completed. Total gradients computed: {len(self.Map_grad)}")

def compute_rmse(warped_pc, target_pc, bidirectional=True, remove_outliers=True):
    """
    计算变形后的源点云与目标点云之间的均方根误差 (RMSE)。

    :param warped_pc: 变形后的源点云 (Open3D PointCloud)
    :param target_pc: 目标点云 (Open3D PointCloud)
    :param bidirectional: 是否进行双向 RMSE 计算
    :param remove_outliers: 是否移除异常点
    :return: RMSE 值 (浮点数)，以及详细的误差统计信息（字典）
    """
    warped_points = np.asarray(warped_pc.points)
    target_points = np.asarray(target_pc.points)
    
    # 初始化最近邻搜索器
    nbrs_target = NearestNeighbors(n_neighbors=1, algorithm='auto').fit(target_points)
    nbrs_warped = NearestNeighbors(n_neighbors=1, algorithm='auto').fit(warped_points)
    
    # 查找每个变形点的最近邻点
    distances_warped, _ = nbrs_target.kneighbors(warped_points)
    # 查找每个目标点的最近邻点（双向）
    if bidirectional:
        distances_target, _ = nbrs_warped.kneighbors(target_points)
    
    # 将距离展平成一维数组
    distances_warped = distances_warped.flatten()
    if bidirectional:
        distances_target = distances_target.flatten()
    
    # 去除异常值
    if remove_outliers:
        # 使用中位数绝对偏差 (MAD) 方法识别异常值
        median_warped = np.median(distances_warped)
        mad_warped = np.median(np.abs(distances_warped - median_warped))
        threshold_warped = median_warped + 3 * mad_warped
        mask_warped = distances_warped <= threshold_warped
        distances_warped = distances_warped[mask_warped]
        
        if bidirectional:
            median_target = np.median(distances_target)
            mad_target = np.median(np.abs(distances_target - median_target))
            threshold_target = median_target + 3 * mad_target
            mask_target = distances_target <= threshold_target
            distances_target = distances_target[mask_target]
    
    # 计算 RMSE
    if bidirectional:
        rmse = np.sqrt((np.mean(distances_warped ** 2) + np.mean(distances_target ** 2)) / 2)
    else:
        rmse = np.sqrt(np.mean(distances_warped ** 2))
    
    # 统计误差信息
    error_stats = {
        'min_error': np.min(distances_warped) if not bidirectional else min(np.min(distances_warped), np.min(distances_target)),
        'max_error': np.max(distances_warped) if not bidirectional else max(np.max(distances_warped), np.max(distances_target)),
        'mean_error': np.mean(distances_warped) if not bidirectional else (np.mean(distances_warped) + np.mean(distances_target)) / 2,
        'median_error': np.median(distances_warped) if not bidirectional else (np.median(distances_warped) + np.median(distances_target)) / 2,
        'std_error': np.std(distances_warped) if not bidirectional else np.sqrt((np.var(distances_warped) + np.var(distances_target)) / 2)
    }
    
    return rmse, error_stats

def compute_mae(warped_pc, target_pc, bidirectional=True, remove_outliers=True):
    """
    计算变形后的源点云与目标点云之间的平均绝对误差 (MAE)。

    :param warped_pc: 变形后的源点云 (Open3D PointCloud)
    :param target_pc: 目标点云 (Open3D PointCloud)
    :param bidirectional: 是否进行双向 MAE 计算
    :param remove_outliers: 是否移除异常点
    :return: MAE 值 (浮点数)，以及详细的误差统计信息（字典）
    """
    warped_points = np.asarray(warped_pc.points)
    target_points = np.asarray(target_pc.points)
    
    # 初始化最近邻搜索器
    nbrs_target = NearestNeighbors(n_neighbors=1, algorithm='auto').fit(target_points)
    nbrs_warped = NearestNeighbors(n_neighbors=1, algorithm='auto').fit(warped_points)
    
    # 查找每个变形点的最近邻点
    distances_warped, _ = nbrs_target.kneighbors(warped_points)
    # 查找每个目标点的最近邻点（双向）
    if bidirectional:
        distances_target, _ = nbrs_warped.kneighbors(target_points)
    
    # 将距离展平成一维数组
    distances_warped = distances_warped.flatten()
    if bidirectional:
        distances_target = distances_target.flatten()
    
    # 去除异常值
    if remove_outliers:
        # 使用中位数绝对偏差 (MAD) 方法识别异常值
        median_warped = np.median(distances_warped)
        mad_warped = np.median(np.abs(distances_warped - median_warped))
        threshold_warped = median_warped + 3 * mad_warped
        mask_warped = distances_warped <= threshold_warped
        distances_warped = distances_warped[mask_warped]
        
        if bidirectional:
            median_target = np.median(distances_target)
            mad_target = np.median(np.abs(distances_target - median_target))
            threshold_target = median_target + 3 * mad_target
            mask_target = distances_target <= threshold_target
            distances_target = distances_target[mask_target]
    
    # 计算 MAE
    if bidirectional:
        mae = (np.mean(np.abs(distances_warped)) + np.mean(np.abs(distances_target))) / 2
    else:
        mae = np.mean(np.abs(distances_warped))
    
    # 统计误差信息
    error_stats = {
        'min_error': np.min(distances_warped) if not bidirectional else min(np.min(distances_warped), np.min(distances_target)),
        'max_error': np.max(distances_warped) if not bidirectional else max(np.max(distances_warped), np.max(distances_target)),
        'mean_error': np.mean(distances_warped) if not bidirectional else (np.mean(distances_warped) + np.mean(distances_target)) / 2,
        'median_error': np.median(distances_warped) if not bidirectional else (np.median(distances_warped) + np.median(distances_target)) / 2,
        'std_error': np.std(distances_warped) if not bidirectional else np.sqrt((np.var(distances_warped) + np.var(distances_target)) / 2)
    }
    
    return mae, error_stats

def homomorphic_mapping(source_pc, target_pc, gridsize=np.pi/250):
    """
    对源点云进行同胚映射，使其在极坐标下的r值与目标DEM的高程值尽可能接近。
    具体来说，调整源点云的每个点，使其与DEM原点的距离（r）等于DEM的高程值（e）。

    :param source_pc: 源点云 (Open3D PointCloud)
    :param target_pc: 目标点云 (Open3D PointCloud)
    :param gridsize: 网格分辨率
    :return: 形变后的源点云 (Open3D PointCloud)
    """
    print("Starting homomorphic mapping...")

    # 1. 计算目标点云的几何中心
    target_center = target_pc.get_center()
    print(f"Target point cloud center: {target_center}")

    # 2. 设置DEM的原点为目标点云的中心
    target_dem = DEM(origin=target_center)  # DEM 原点设置为目标中心
    target_dem.DEM_generate_polar(target_pc, gridsize)
    target_map_elevation = target_dem.GetMap_e()

    # 输出 DEM 高程值的统计信息
    dem_elevations = np.array(list(target_map_elevation.values()))
    print(f"DEM Elevation statistics: min={dem_elevations.min():.3f}, max={dem_elevations.max():.3f}, mean={dem_elevations.mean():.3f}, std={dem_elevations.std():.3f}")

    # 3. 将源点云转换为极坐标系
    points = np.asarray(source_pc.points)
    colors = np.asarray(source_pc.colors)
    # 以DEM的原点为参考点进行计算，不平移点云
    vectors = points - target_center  # 计算每个点相对于DEM原点的向量
    r = np.linalg.norm(vectors, axis=1)
    # 避免除以零
    with np.errstate(divide='ignore', invalid='ignore'):
        directions = np.divide(vectors, r[:, np.newaxis])
        directions[~np.isfinite(directions)] = 0  # 替换无穷大和NaN为0

    # 计算theta和phi
    theta = np.arctan2(vectors[:, 1], vectors[:, 0])
    theta[theta < 0] += 2 * np.pi
    phi = np.arccos(np.clip(vectors[:, 2] / r, -1.0, 1.0))
    phi[r == 0] = 0

    # 计算网格索引
    nrow = int(np.ceil(target_dem.width_ / gridsize))
    ncol = int(np.ceil(target_dem.height_ / gridsize))
    rowIDs = np.floor(theta / gridsize).astype(int)
    colIDs = np.floor(phi / gridsize).astype(int)

    # 过滤超出范围的索引
    rowIDs = np.clip(rowIDs, 0, nrow - 1)
    colIDs = np.clip(colIDs, 0, ncol - 1)

    # 为 DEM 网格创建 theta 和 phi 坐标
    theta_grid = []
    phi_grid = []
    elevation_grid = []
    for (row, col), elev in target_map_elevation.items():
        theta_center = row * gridsize + gridsize / 2
        phi_center = col * gridsize + gridsize / 2
        theta_grid.append(theta_center)
        phi_grid.append(phi_center)
        elevation_grid.append(elev)
    points_grid = np.vstack((theta_grid, phi_grid)).T
    elevation_grid = np.array(elevation_grid)

    # 确保 points_grid 和 elevation_grid 的长度一致
    assert points_grid.shape[0] == elevation_grid.shape[0], "Points grid and elevation grid must have the same number of elements."

    # 使用线性插值填补 DEM 缺失区域
    interpolated_elevations = griddata(
        points_grid, 
        elevation_grid, 
        (theta, phi), 
        method='linear'
    )

    # 对于仍有缺失值的点，使用最近邻插值填补
    missing = np.isnan(interpolated_elevations)
    if np.any(missing):
        interpolated_elevations[missing] = griddata(
            points_grid, 
            elevation_grid, 
            (theta[missing], phi[missing]), 
            method='nearest'
        )

    # 检查插值后的高程值
    print(f"Interpolated Elevation statistics: min={interpolated_elevations.min():.3f}, max={interpolated_elevations.max():.3f}, mean={interpolated_elevations.mean():.3f}, std={interpolated_elevations.std():.3f}")

    # 4. 设置新的 r 值为插值后的 DEM 高程值
    new_r = interpolated_elevations

    # 5. 处理 r == 0 的点（避免将所有这些点移动到原点）
    zero_r_mask = (r == 0)
    if np.any(zero_r_mask):
        # 为这些点分配一个默认方向向量，例如 [1, 0, 0]
        directions[zero_r_mask] = np.array([1, 0, 0])

    # 6. 计算新的点坐标
    adjusted_cartesian_points = target_center + directions * new_r[:, np.newaxis]

    # 创建变形后的点云
    warped_pc = o3d.geometry.PointCloud()
    warped_pc.points = o3d.utility.Vector3dVector(adjusted_cartesian_points)
    warped_pc.colors = o3d.utility.Vector3dVector(colors)

    # 统计调整的点数
    adjusted_count = len(points)  # 所有点都被调整
    print(f"Adjusted {adjusted_count} out of {len(points)} points.")
    print(f"Percentage of points adjusted: {adjusted_count / len(points) * 100:.2f}%")

    # 统计调整前后的 r 值差异
    original_r = r
    adjusted_r = np.linalg.norm(adjusted_cartesian_points - target_center, axis=1)
    r_diff = adjusted_r - original_r
    adjusted_indices = r_diff != 0

    # 计算调整后的 r 值统计
    print("Adjusted r values statistics:")
    if np.any(adjusted_indices):
        print(f"  r difference min: {np.min(r_diff[adjusted_indices]):.3f}")
        print(f"  r difference max: {np.max(r_diff[adjusted_indices]):.3f}")
        print(f"  r difference mean: {np.mean(r_diff[adjusted_indices]):.3f}")
        print(f"  r difference std: {np.std(r_diff[adjusted_indices]):.3f}")
    else:
        print("  No r differences found.")

    print("Homomorphic mapping completed.")
    return warped_pc

def compute_rmse(warped_pc, target_pc, bidirectional=True, remove_outliers=True):
    """
    计算变形后的源点云与目标点云之间的均方根误差 (RMSE)。

    :param warped_pc: 变形后的源点云 (Open3D PointCloud)
    :param target_pc: 目标点云 (Open3D PointCloud)
    :param bidirectional: 是否进行双向 RMSE 计算
    :param remove_outliers: 是否移除异常点
    :return: RMSE 值 (浮点数)，以及详细的误差统计信息（字典）
    """
    warped_points = np.asarray(warped_pc.points)
    target_points = np.asarray(target_pc.points)
    
    # 初始化最近邻搜索器
    nbrs_target = NearestNeighbors(n_neighbors=1, algorithm='auto').fit(target_points)
    nbrs_warped = NearestNeighbors(n_neighbors=1, algorithm='auto').fit(warped_points)
    
    # 查找每个变形点的最近邻点
    distances_warped, _ = nbrs_target.kneighbors(warped_points)
    # 查找每个目标点的最近邻点（双向）
    if bidirectional:
        distances_target, _ = nbrs_warped.kneighbors(target_points)
    
    # 将距离展平成一维数组
    distances_warped = distances_warped.flatten()
    if bidirectional:
        distances_target = distances_target.flatten()
    
    # 去除异常值
    if remove_outliers:
        # 使用中位数绝对偏差 (MAD) 方法识别异常值
        median_warped = np.median(distances_warped)
        mad_warped = np.median(np.abs(distances_warped - median_warped))
        threshold_warped = median_warped + 3 * mad_warped
        mask_warped = distances_warped <= threshold_warped
        distances_warped = distances_warped[mask_warped]
        
        if bidirectional:
            median_target = np.median(distances_target)
            mad_target = np.median(np.abs(distances_target - median_target))
            threshold_target = median_target + 3 * mad_target
            mask_target = distances_target <= threshold_target
            distances_target = distances_target[mask_target]
    
    # 计算 RMSE
    if bidirectional:
        rmse = np.sqrt((np.mean(distances_warped ** 2) + np.mean(distances_target ** 2)) / 2)
    else:
        rmse = np.sqrt(np.mean(distances_warped ** 2))
    
    # 统计误差信息
    error_stats = {
        'min_error': np.min(distances_warped) if not bidirectional else min(np.min(distances_warped), np.min(distances_target)),
        'max_error': np.max(distances_warped) if not bidirectional else max(np.max(distances_warped), np.max(distances_target)),
        'mean_error': np.mean(distances_warped) if not bidirectional else (np.mean(distances_warped) + np.mean(distances_target)) / 2,
        'median_error': np.median(distances_warped) if not bidirectional else (np.median(distances_warped) + np.median(distances_target)) / 2,
        'std_error': np.std(distances_warped) if not bidirectional else np.sqrt((np.var(distances_warped) + np.var(distances_target)) / 2)
    }
    
    return rmse, error_stats

def compute_mae(warped_pc, target_pc, bidirectional=True, remove_outliers=True):
    """
    计算变形后的源点云与目标点云之间的平均绝对误差 (MAE)。

    :param warped_pc: 变形后的源点云 (Open3D PointCloud)
    :param target_pc: 目标点云 (Open3D PointCloud)
    :param bidirectional: 是否进行双向 MAE 计算
    :param remove_outliers: 是否移除异常点
    :return: MAE 值 (浮点数)，以及详细的误差统计信息（字典）
    """
    warped_points = np.asarray(warped_pc.points)
    target_points = np.asarray(target_pc.points)
    
    # 初始化最近邻搜索器
    nbrs_target = NearestNeighbors(n_neighbors=1, algorithm='auto').fit(target_points)
    nbrs_warped = NearestNeighbors(n_neighbors=1, algorithm='auto').fit(warped_points)
    
    # 查找每个变形点的最近邻点
    distances_warped, _ = nbrs_target.kneighbors(warped_points)
    # 查找每个目标点的最近邻点（双向）
    if bidirectional:
        distances_target, _ = nbrs_warped.kneighbors(target_points)
    
    # 将距离展平成一维数组
    distances_warped = distances_warped.flatten()
    if bidirectional:
        distances_target = distances_target.flatten()
    
    # 去除异常值
    if remove_outliers:
        # 使用中位数绝对偏差 (MAD) 方法识别异常值
        median_warped = np.median(distances_warped)
        mad_warped = np.median(np.abs(distances_warped - median_warped))
        threshold_warped = median_warped + 3 * mad_warped
        mask_warped = distances_warped <= threshold_warped
        distances_warped = distances_warped[mask_warped]
        
        if bidirectional:
            median_target = np.median(distances_target)
            mad_target = np.median(np.abs(distances_target - median_target))
            threshold_target = median_target + 3 * mad_target
            mask_target = distances_target <= threshold_target
            distances_target = distances_target[mask_target]
    
    # 计算 MAE
    if bidirectional:
        mae = (np.mean(np.abs(distances_warped)) + np.mean(np.abs(distances_target))) / 2
    else:
        mae = np.mean(np.abs(distances_warped))
    
    # 统计误差信息
    error_stats = {
        'min_error': np.min(distances_warped) if not bidirectional else min(np.min(distances_warped), np.min(distances_target)),
        'max_error': np.max(distances_warped) if not bidirectional else max(np.max(distances_warped), np.max(distances_target)),
        'mean_error': np.mean(distances_warped) if not bidirectional else (np.mean(distances_warped) + np.mean(distances_target)) / 2,
        'median_error': np.median(distances_warped) if not bidirectional else (np.median(distances_warped) + np.median(distances_target)) / 2,
        'std_error': np.std(distances_warped) if not bidirectional else np.sqrt((np.var(distances_warped) + np.var(distances_target)) / 2)
    }
    
    return mae, error_stats

def visualize_error_distribution(warped_pc, target_pc):
    """
    可视化变形后的源点云与目标点云之间的误差分布。

    :param warped_pc: 变形后的源点云 (Open3D PointCloud)
    :param target_pc: 目标点云 (Open3D PointCloud)
    """
    warped_points = np.asarray(warped_pc.points)
    target_points = np.asarray(target_pc.points)
    
    # 初始化最近邻搜索器
    nbrs_target = NearestNeighbors(n_neighbors=1, algorithm='auto').fit(target_points)
    distances_warped, _ = nbrs_target.kneighbors(warped_points)
    distances_warped = distances_warped.flatten()
    
    # 使用中位数绝对偏差 (MAD) 方法识别并去除异常值
    median = np.median(distances_warped)
    mad = np.median(np.abs(distances_warped - median))
    threshold = median + 3 * mad
    filtered_distances = distances_warped[distances_warped <= threshold]
    
    # 绘制误差分布直方图
    plt.figure(figsize=(10, 6))
    plt.hist(filtered_distances, bins=50, color='skyblue', edgecolor='black')
    plt.title('Error Distribution (RMSE Evaluation)')
    plt.xlabel('Distance')
    plt.ylabel('Frequency')
    plt.grid(True)
    plt.show()

def compute_r2(warped_pc, target_pc):
    """
    计算变形后的源点云与目标点云之间的 R² 决定系数。

    :param warped_pc: 变形后的源点云 (Open3D PointCloud)
    :param target_pc: 目标点云 (Open3D PointCloud)
    :return: R² 值 (浮点数)
    """
    warped_points = np.asarray(warped_pc.points)
    target_points = np.asarray(target_pc.points)
    
    # 初始化最近邻搜索器
    nbrs_target = NearestNeighbors(n_neighbors=1, algorithm='auto').fit(target_points)
    
    # 查找每个变形点的最近邻点
    distances, indices = nbrs_target.kneighbors(warped_points)
    nearest_target_points = target_points[indices.flatten()]
    
    # 计算总平方和 (Total Sum of Squares)
    ss_tot = np.sum((nearest_target_points - np.mean(nearest_target_points, axis=0)) ** 2)
    
    # 计算残差平方和 (Residual Sum of Squares)
    ss_res = np.sum((nearest_target_points - warped_points) ** 2)
    
    # 计算 R²
    r2 = 1 - (ss_res / ss_tot)
    
    return r2

def save_point_cloud(pc, file_path, write_ascii=True):
    """
    保存点云到指定文件路径。

    :param pc: 点云对象 (Open3D PointCloud)
    :param file_path: 保存的文件路径 (字符串)
    :param write_ascii: 是否以 ASCII 格式保存 (布尔值)
    """
    try:
        o3d.io.write_point_cloud(file_path, pc, write_ascii=write_ascii)
        print(f"Point cloud successfully saved to {file_path}")
    except Exception as e:
        print(f"保存点云时出错: {e}")

def visualize_dem_coverage(dem, gridsize):
    """
    可视化 DEM 覆盖范围。

    :param dem: DEM 对象
    :param gridsize: 网格分辨率
    """
    theta_ids = [key[0] for key in dem.GetMap_e().keys()]
    phi_ids = [key[1] for key in dem.GetMap_e().keys()]
    elevations = list(dem.GetMap_e().values())

    plt.figure(figsize=(12, 6))
    plt.scatter(theta_ids, phi_ids, c=elevations, cmap='viridis', s=1)
    plt.xlabel('Theta Grid ID')
    plt.ylabel('Phi Grid ID')
    plt.title('DEM Coverage Visualization')
    plt.colorbar(label='Elevation')
    plt.show()

def main():
    """
    主函数，执行同胚映射过程并评估精度。
    """
    # 定义源点云和目标点云的路径
    target_pc_path = "/home/lyt/LCode/point_test/aligne11.ply"  # 替换为实际路径
    source_pc_path = "/home/lyt/LCode/point_test/dummy1.ply"
    #source_pc_path = "/home/lyt/LCode/new_point/Xa_prime_with_normals25.ply"  # 替换为实际路径

    # 确保保存目录存在
    save_dir = "/home/lyt/LCode/new_point/"
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    try:
        print(f"Loading source point cloud from {source_pc_path}...")
        source_pc = o3d.io.read_point_cloud(source_pc_path)
        print(f"Source point cloud loaded: {len(source_pc.points)} points.")
    
        print(f"Loading target point cloud from {target_pc_path}...")
        target_pc = o3d.io.read_point_cloud(target_pc_path)
        print(f"Target point cloud loaded: {len(target_pc.points)} points.")
    except Exception as e:
        print(f"加载点云时出错: {e}")
        return

    # 检查点云是否包含颜色信息
    if not source_pc.has_colors():
        print("源点云没有颜色信息，初始化为红色。")
        source_pc.paint_uniform_color([1, 0, 0])  # 红色
    else:
        print("源点云包含颜色信息。")

    if not target_pc.has_colors():
        print("目标点云没有颜色信息，初始化为绿色。")
        target_pc.paint_uniform_color([0, 1, 0])  # 绿色
    else:
        print("目标点云包含颜色信息。")

    # 进行同胚映射
    print("Performing homomorphic mapping...")
    warped_source_pc = homomorphic_mapping(source_pc, target_pc, gridsize=np.pi/250)  # 调整 gridsize

    # 保存变形后的点云
    warped_pc_save_path = os.path.join(save_dir, "1.ply")  # 替换为您希望保存的路径
    save_point_cloud(warped_source_pc, warped_pc_save_path, write_ascii=True)

    # 计算并输出 RMSE
    print("Computing RMSE between warped source and target point clouds...")
    rmse, error_stats = compute_rmse(warped_source_pc, target_pc, bidirectional=True, remove_outliers=True)
    print(f"RMSE between warped source and target point clouds: {rmse:.3f}")
    print("Error Statistics:")
    for key, value in error_stats.items():
        print(f"  {key}: {value:.3f}")

    # 计算并输出 MAE
    print("Computing MAE between warped source and target point clouds...")
    mae, mae_stats = compute_mae(warped_source_pc, target_pc, bidirectional=True, remove_outliers=True)
    print(f"MAE between warped source and target point clouds: {mae:.3f}")
    print("MAE Statistics:")
    for key, value in mae_stats.items():
        print(f"  {key}: {value:.3f}")

    # 计算并输出 R² 决定系数
    print("Computing R² between warped source and target point clouds...")
    r2 = compute_r2(warped_source_pc, target_pc)
    print(f"R² between warped source and target point clouds: {r2:.3f}")

    # 可视化误差分布
    visualize_error_distribution(warped_pc=warped_source_pc, target_pc=target_pc)

    # 可视化 DEM 覆盖范围
    target_dem = DEM(origin=target_pc.get_center())
    target_dem.DEM_generate_polar(target_pc, gridsize=np.pi/250)
    visualize_dem_coverage(target_dem, gridsize=np.pi/250)

    # 可视化原始源点云、目标 DEM 和形变后的源点云
    # 为了更好地可视化，建议使用不同的颜色
    source_pc_temp = source_pc.paint_uniform_color([1, 0, 0])  # 红色表示源点云
    target_pc_temp = target_pc.paint_uniform_color([0, 1, 0])  # 绿色表示目标点云
    warped_source_pc_temp = warped_source_pc.paint_uniform_color([0, 0, 1])  # 蓝色表示形变后的源点云

    # 添加坐标轴以便参考
    # 坐标轴的原点设置为DEM的原点（目标点云的几何中心）
    target_center = target_pc.get_center()
    axis_pcd = o3d.geometry.TriangleMesh.create_coordinate_frame(size=1.0, origin=target_center)

    # 可视化
    print("Visualizing point clouds...")
    o3d.visualization.draw_geometries([
        source_pc_temp, 
        target_pc_temp, 
        warped_source_pc_temp,
        axis_pcd
    ],
        window_name="Homomorphic Mapping",
        width=800, 
        height=600
    )

if __name__ == "__main__":
    main()
